#!/usr/bin/env python

""" face_detector.py - Version 0.1 2011-06-08

    Implements a Haar face detector on an RGB video stream using OpenCV and the pi_rein
    detector.py class.

    Created for the Pi Robot Project: http://www.pirobot.org
    Copyright (c) 2011 Patrick Goebel.  All rights reserved.

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.5
    
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details at:
    
    http://www.gnu.org/licenses/gpl.html
      
"""

import roslib
roslib.load_manifest('pi_rein')
from detector import Detector
import rospy
import cv
import sys
from pi_rein.msg import *
from pi_rein.srv import *
from sensor_msgs.msg import Image, PointCloud2
from std_msgs.msg import Header, ColorRGBA
import random
import os

class FaceDetector(Detector):
    def __init__(self, node_name):
        Detector.__init__(self, node_name)
        
        self.face_pub = rospy.Publisher("detections", DetectionArray)
            
        self.grey = None
        self.small_image = None
        
        self.rate = rospy.Rate(rospy.get_param("~rate", 10))
        
        self.cascade_frontal_alt = rospy.get_param("~cascade_frontal_alt", "")
        self.cascade_frontal_alt2 = rospy.get_param("~cascade_frontal_alt2", "")
        self.cascade_profile = rospy.get_param("~cascade_profile", "")
        
        self.cascade_frontal_alt = cv.Load(self.cascade_frontal_alt)
        self.cascade_frontal_alt2 = cv.Load(self.cascade_frontal_alt2)
        self.cascade_profile = cv.Load(self.cascade_profile)
        
        self.use_masks = rospy.get_param("~use_masks", False)
        
        self.min_size = (20, 20)
        self.image_scale = 2
        self.haar_scale = 1.5
        self.min_neighbors = 1
        self.haar_flags = cv.CV_HAAR_DO_CANNY_PRUNING

        self.faces = DetectionArray()
        self.max_faces = 2
        self.facebox_color = ColorRGBA(255, 0, 0, 0)
        
        rospy.wait_for_message("image", Image, timeout=15)

    def image_callback(self, ros_image):
        (self.image, self.image_size, self.image_width, self.image_height) = self.ros2cv_image(ros_image, "bgr8")
        self.image_stamp = ros_image.header.stamp
        try:
            self.detect()
            self.faces = self.getDetections()
            #if len(self.faces.detections) > 0:
            self.face_pub.publish(self.faces)
            if self.use_display:
                self.draw_markers(self.faces.header, self.faces.detections)
        except:
            pass
                    
    def detect(self):
        faces = None
        
        """ Clear any left over detections """
        self.detections = DetectionArray()
        
        if not self.grey:
            """ Allocate temporary images """      
            self.grey = cv.CreateImage(self.image_size, 8, 1)
            self.small_image = cv.CreateImage((cv.Round(self.image_width / self.image_scale),
                       cv.Round(self.image_height / self.image_scale)), 8, 1)
    
        """ Convert color input image to grayscale """
        cv.CvtColor(self.image, self.grey, cv.CV_BGR2GRAY)
        
        """ Equalize the histogram to reduce lighting effects. """
        cv.EqualizeHist(self.grey, self.grey)
        
        #cv.SetImageROI(self.grey, self.masks.masks[0].roi)
            
        """ Scale input image for faster processing """
        cv.Resize(self.grey, self.small_image, cv.CV_INTER_LINEAR)
    
        """ First check one of the frontal templates """
        if self.cascade_frontal_alt2:
            faces = cv.HaarDetectObjects(self.small_image, self.cascade_frontal_alt2, cv.CreateMemStorage(0),
                                          self.haar_scale, self.min_neighbors, self.haar_flags, self.min_size)
                                         
        """ If that fails, check the profile template """
        if not faces:
            if self.cascade_profile:
                faces = cv.HaarDetectObjects(self.small_image, self.cascade_profile, cv.CreateMemStorage(0),
                                             self.haar_scale, self.min_neighbors, self.haar_flags, self.min_size)

            if not faces:
                """ If that fails, check a different frontal profile """
                if self.cascade_frontal_alt:
                    faces = cv.HaarDetectObjects(self.small_image, self.cascade_frontal_alt, cv.CreateMemStorage(0),
                                         self.haar_scale, self.min_neighbors, self.haar_flags, self.min_size)
        
        #cv.ResetImageROI(self.grey)
        n_faces = len(faces)
                        
        for i in range(min(n_faces, self.max_faces)):
            ((x, y, w, h), n) = faces[i]
            """ The input to cv.HaarDetectObjects was resized, so scale the 
                bounding box of each face and convert it to two CvPoints """
            pt1 = (int(x * self.image_scale), int(y * self.image_scale))
            pt2 = (int((x + w) * self.image_scale), int((y + h) * self.image_scale))
            face_width = pt2[0] - pt1[0]
            face_height = pt2[1] - pt1[1]
            
            detection = Detection()
            detection.label = 'face_' + str(i)
            detection.detector = 'face_' + str(i)
            detection.header = Header()
            detection.mask = Mask()
            detection.mask.roi = Rect(pt1[0], pt1[1], face_width, face_height)
            detection.header.stamp = self.image_stamp
            detection.header.frame_id = '/openni_rgb_optical_frame'
            if i == 0:
                detection.marker_color = ColorRGBA(0, 255, 0, 0)
            else:
                detection.marker_color = self.facebox_color

            self.detections.detections.append(detection)
            
        self.detections.header.stamp = self.image_stamp
        self.detections.header.frame_id = '/openni_rgb_optical_frame'
    
    def setImage(self, image):
        pass
    
    def setPointCloud(self, point_cloud):
        pass
    
    def loadModels(self):
        pass      

def main(args):
    # Display a help message if appropriate.
    help_message = ""
          
    print help_message
    
    try:
        FaceDetector("haar_face_detector")
        rospy.spin()
    except KeyboardInterrupt:
        print "Shutting down detector node."

if __name__ == '__main__':
    main(sys.argv)